from typing import Dict, List, Tuple, Optional
import rasterio
import numpy as np
import os
from calc_v1.utils import load_parameters_from_json, extract_codes


def split_tif_by_ecosystem(
    eco_tif_path: str,
    gep_res_tif_path: str,
    output_dir: str = None,
    eco_type_json_path: str = None,
    compress: str = "lzw",
    bigtiff: str = "YES"
) -> Tuple[bool, Optional[str]]:
    """
    根据生态图分类拆分空气质量核算 TIFF 文件。

    :param eco_tif_path: 生态图 TIFF 文件路径
    :param gep_res_tif_path: 空气质量/土壤保持等核算结果 TIFF 文件路径
    :param output_dir: 输出目录路径
    :param eco_type_json_path: 包含生态系统映射的 JSON 文件路径
    :param compress: TIFF 压缩方式，默认 'lzw'
    :param bigtiff: 是否启用 BIGTIFF 支持，默认 'YES'

    :return: (成功标志, 错误信息)
    """
    if output_dir is None:
        output_dir = os.path.join(os.path.dirname(gep_res_tif_path), "spiltByEcoOutput")
    if eco_type_json_path is None:
        eco_type_json_path = r"F:\code\dev\calc-gep-regulate-cqc\calc_v1\data\ecosystems_json\ecosystems_type_GB.json"

    try:
        # 读取生态系统映射
        ecosystem_json = load_parameters_from_json(eco_type_json_path)
        eco_mapping: Dict[str, List[int]] = {}
        for ecosystem in ecosystem_json["applicable_ecosystems"]:
            name = ecosystem["name"]
            codes = extract_codes(ecosystem)
            eco_mapping[name] = codes

        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)

        # 读取生态图数据
        with rasterio.open(eco_tif_path) as eco_src:
            eco_data: np.ndarray = eco_src.read(1)
            eco_nodata: float = eco_src.nodata
            eco_profile = eco_src.profile
            eco_transform = eco_src.transform
            eco_crs = eco_src.crs

        # 读取空气质量数据
        with rasterio.open(gep_res_tif_path) as air_src:
            air_data: np.ndarray = air_src.read(1)
            air_nodata: float = air_src.nodata

        # 检查分辨率是否一致
        assert eco_data.shape == air_data.shape, "生态图与空气质量图尺寸不一致！请确保它们的空间对齐。"

        # 创建掩膜
        if eco_nodata is not None:
            eco_valid_mask = eco_data != eco_nodata
        else:
            eco_valid_mask = np.ones_like(eco_data, dtype=bool)

        if air_nodata is not None:
            air_valid_mask = air_data != air_nodata
        else:
            air_valid_mask = np.ones_like(air_data, dtype=bool)

        valid_mask = eco_valid_mask & air_valid_mask

        # 构建基础文件名
        base_filename = os.path.basename(gep_res_tif_path)
        parts = base_filename.split("_")

        # 遍历每个生态系统类型
        for eco_name, eco_codes in eco_mapping.items():
            eco_mask = np.isin(eco_data, eco_codes)
            final_mask = eco_mask & valid_mask
            extracted_data = np.where(final_mask, air_data, np.nan)

            # 如果全是 NaN，则跳过
            if np.isnan(extracted_data).all():
                print(f"{eco_name} 在生态图中无对应区域，跳过保存。")
                continue

            # 替换文件名第三段字段
            if len(parts) > 2:
                new_parts = parts[:2] + [eco_name] + parts[3:]
                new_filename = "_".join(new_parts)
            else:
                raise ValueError("文件名格式不符合预期，无法找到第3个字段")

            output_path = os.path.join(output_dir, new_filename)

            # 更新 profile
            eco_profile.update(
                dtype=rasterio.float32,
                nodata=air_nodata,
                count=1,
                crs=eco_crs,
                transform=eco_transform,
                driver='GTiff',
                compress=compress,
                bigtiff=bigtiff
            )

            # 写入文件
            with rasterio.open(output_path, 'w', **eco_profile) as dst:
                dst.write(extracted_data.astype(np.float32), 1)

            print(f"已保存 {eco_name} 的核算结果到 {output_path}")

        return True, None

    except Exception as e:
        print(f"发生错误：{e}")
        return False, str(e)

if __name__ == "__main__":
    success, error = split_tif_by_ecosystem(
        eco_tif_path=r"H:\test\eco_classify\china_110000.tif",
        gep_res_tif_path=r"H:\test\eco_classify\110000_Qsr_Qsr_replaced_output.tif",
        output_dir=r"H:\test\eco_classify\output",
    )

    if success:
        print("拆分完成！")
    else:
        print(f"拆分失败：{error}")
